/*
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

#ifndef LIB_RTTI_H_
#define LIB_RTTI_H_

#include <algorithm>
#include <array>
#include <cstdint>
#include <type_traits>

/// @file
/// Custom intrusive lightweight RTTI implementation for semi-open class hierarchies.
///
/// P4C IR::Node class herarchy relies a lot on virtual bases and at the same
/// time does a lot of dynamic type checks and downcasts /
/// upcasts. Unfortunately, `static_cast` cannot be used to cast to / via
/// virtual bases and use of `dynamic_cast` involves a non-negligible overhead.
/// The typical IR::Node class hierarchy looks like this:
///   class Node : public virtual INode;
///   class StatOrDecl : public Node;
///   class IDeclaration : public virtual INode;
///   class Declaration : public StatOrDecl, public virtual IDeclaration;
/// and it is pretty common code pattern to do something like this:
///   Node *n = new Declaration();
///   n->to<IDeclaration>();
///
/// The proposed solution relies on intrusive RTTI metadata that is either autogenerated
/// (via ir-generator for IR::Node and friends) or provided explicitly by anyone willing
/// to use lightweight RTTI. Two main components of RTTI are:
///  - 64-bit typeids that could be either synthesized from class name at
///    compile time or explicitly specified
///  - Set of functions performing dynamic type checks and class hierarchy traversal
///
///  For example above in order to perform downcast to a virtual base we are
///  doing "down-and-up" trick using virtual dispatch for `this` pointer adjustment. Briefly:
///    - Inside `to<T>` we determine the typeid of IDeclaration
///    - Call virtually toImpl, it lands with proper `this` inside
///    - Declaration::toImpl; we know all immediate bases of Declaration via
///      RTTI::TypeInfo<> due to RTTI annotation
///    - From there we do upcast via TypeInfo::dyn_cast, finding the base type that corresponds
///      to a typeId (of IDeclaration).
///
///  The implementation itself was crafted to allow as much compiler
///  optimizations as possible. For example `is<T> / isA` after inlining
///  compiles down to several integer comparisons (that are often combined into
///  range checks). TypeInfo::dyn_cast looks expensive, however, everything is
///  inlined and therefore is again series of typeid ranged checks interleaved
///  with `this` adjustment.
///  The code relies on C++17 features. Care was taken to ensure that `constexpr`
///  functions are compiled down to compile-time constants for GCC >= 9 and recent
///  of clang. Some of them likely would need to be marked as `consteval` when C++20
///  would be used across the codebase.

namespace P4::RTTI {

using TypeId = std::uint64_t;
static constexpr TypeId InvalidTypeId = UINT64_C(0);
static constexpr unsigned kDiscriminatorBits = 56;
static constexpr uint64_t kInnerTypeIdMask = (UINT64_C(1) << kDiscriminatorBits) - 1;
static constexpr uint64_t kHashDiscriminator = UINT64_C(0xFF);

namespace detail {
// Apparently string_view does not optimize properly in constexpr context
// for GCC < 13
struct TypeNameHolder {
    template <size_t N>
    constexpr TypeNameHolder(const char (&str)[N]) : str(&str[0]), length(N - 1) {}  // NOLINT

    const char *str;
    size_t length;
};

// Derive the string containing the ultimate typename of `T`. The resulting
// string should be available in constexpr context.
// TODO: Probably use std::source_location::function_name with C++20
template <typename T>
constexpr TypeNameHolder typeNameProxy() {
#ifdef __clang__
    return __PRETTY_FUNCTION__;
#elif defined(__GNUC__)
    return __PRETTY_FUNCTION__;
#else
// TODO: We might use __FUNCSIG__ with MSVC should this will be necessary.
#error "Unsupported compiler"
#endif
}

static constexpr uint64_t FNV1a(const char *str, size_t n,
                                uint64_t hash = UINT64_C(0xcbf29ce484222325)) {
    return n == 0 ? hash : FNV1a(str + 1, n - 1, UINT64_C(0x100000001b3) * (hash ^ str[0]));
}

static constexpr uint64_t FNV1a(TypeNameHolder str) { return FNV1a(str.str, str.length); }

// TypeIdResolver provides a unified way to getting typeid of the type `T`. Typeid
// could be derived in two different ways:
//   * Explicitly, if `T` implements a `T::static_typeId()` method that could be
//     used in a constexpr context
//   * Implicitly, via hash from typename itself
//  Also, the explicit specification could fallback to an implicit one, of the
//  returned typeid is `InvalidTypeId`
//
//  The top 8 bits of typeid are reserved as discriminator and could be used to
//  delineate different typeids. Currently the following discriminators are reserved:
//   - 0xFF is used to mark all hash-derived typeids
//   - 0x00 is used to mark all explicitly-specified typeids
//  In addition, IR::Node hierarchy uses the following discriminators:
//   - 0x01 for Vector<T> with low 56 bits contain typeid of `T`.
//   - 0x02 for IndexedVector<T> with low 56 bits contain typeid of `T`.
//  One could use DECLARE_TYPEINFO_WITH_NESTED_TYPEID to combine typeids in such way.

// Default implementation of typeid resolver: use hash derived from type name,
// but mark upper 8 bits to be 0xFF
template <typename T, typename = void>
struct TypeIdResolver {
    static constexpr TypeId resolve() noexcept {
        // Get original 64-bit hash
        uint64_t hash = FNV1a(typeNameProxy<T>());
        // Xor-fold to obtain 56-bit value
        hash = (hash >> kDiscriminatorBits) ^ (hash & kInnerTypeIdMask);
        return (kHashDiscriminator << kDiscriminatorBits) | hash;
    }
};

// Specialized implementation that looks for static_typeId member.
template <typename T>
struct TypeIdResolver<T, std::void_t<decltype(T::static_typeId)>> {
    static constexpr uint64_t resolve() noexcept {
        constexpr TypeId id = T::static_typeId();
        // Allow fall-back implementation if class does not want to explicitly
        // specify a typeid
        if constexpr (id == InvalidTypeId) return TypeIdResolver<T, int>::resolve();

        return id;
    }
};

}  // namespace detail

/// Given a "full" typeid, returns one with discriminator removed
static constexpr TypeId innerTypeId(TypeId id) { return id & kInnerTypeIdMask; }

/// Extracts a discriminator from typeid
static constexpr TypeId typeidDiscriminator(TypeId id) { return id >> kDiscriminatorBits; }

/// Combines typeid and discriminator
static constexpr TypeId combineTypeIdWithDiscriminator(TypeId discriminator, TypeId id) {
    return (discriminator << kDiscriminatorBits) | innerTypeId(id);
}

struct Base;

template <typename This, typename... Parents>
struct TypeInfo {
    using T = std::remove_const_t<std::remove_reference_t<This>>;

    static_assert((... && std::is_base_of_v<Parents, This>),
                  "One or more parents are not a base of this type.");

    static_assert((... && std::is_base_of_v<Base, Parents>),
                  "One or more parent hierarchies is not based on top of RTTI::Base.");

    static_assert(std::is_base_of_v<Base, This>,
                  "Invalid typeinfo requested, class is not based on top of RTTI::Base.");

    static_assert(
        std::is_same_v<T *, decltype(T::rttiEnabledMarker(nullptr))>,
        "Invalid typeinfo requested, class does not declare typeinfo via DECLARE_TYPEINFO.");

    [[nodiscard]] static constexpr TypeId id() noexcept {
        return detail::TypeIdResolver<T>::resolve();
    }

    [[nodiscard]] static constexpr bool isA(TypeId typeId) noexcept {
        return (id() == typeId) || (... || (Parents::TypeInfo::isA(typeId)));
    }

    template <typename T>
    [[nodiscard]] static const void *dyn_cast(TypeId typeId, const T *ptr) noexcept {
        if (id() == typeId) return static_cast<const This *>(ptr);

        std::array<const void *, sizeof...(Parents)> ptrs = {
            Parents::TypeInfo::dyn_cast(typeId, ptr)...};

        auto it =
            std::find_if(ptrs.begin(), ptrs.end(), [](const void *ptr) { return ptr != nullptr; });
        return (it != ptrs.end()) ? *it : nullptr;
    }
};

struct Base {
    virtual ~Base() = default;

    /// Queries information about dynamic type of an object.
    /// @return typeid of current object.
    [[nodiscard]] virtual TypeId typeId() const noexcept = 0;
    /// Checks if given object is of specified dynamic type.
    /// @param typeid to check with
    [[nodiscard]] virtual bool isA(TypeId typeId) const noexcept = 0;

    /// Checks if given object is of specified dynamic type.
    /// Same as `isA`, but typeid to check with is derived from template argument.
    template <typename T>
    [[nodiscard]] bool is() const noexcept {
        return isA(TypeInfo<T>::id());
    }

    /// Casts current object to a given type `T`.
    /// Overall the semantics is similar to that of `dynamic_cast`, however,
    /// relies on custom RTTI metadata and intrusive annotation.
    /// @return nullptr if object is not compatible with `T` (e.g. if `T` is not
    /// among its base classes), casted pointer to object of T* type otherwise.
    template <typename T>
    [[nodiscard]] T *to() noexcept {
        return reinterpret_cast<T *>(const_cast<void *>(toImpl(TypeInfo<T>::id())));
    }

    /// Same as `to`, but returns const pointer to T.
    template <typename T>
    [[nodiscard]] const T *to() const noexcept {
        return reinterpret_cast<const T *>(toImpl(TypeInfo<T>::id()));
    }

 protected:
    [[nodiscard]] virtual const void *toImpl(TypeId typeId) const noexcept = 0;
};

}  // namespace P4::RTTI

/// Declare typeinfo for a given class `T`.
///
/// All direct bases (including virtual) should be specified as additional
/// arguments. For example:
/// struct Base1 {
///   ...
///   DECLARE_TYPEINFO(Base1);
/// };
/// struct Base2 {
///   ...
///   DECLARE_TYPEINFO(Base2);
/// };
///
/// struct Derived : public Base1, public virtual Base2 {
///   ...
///   DECLARE_TYPEINFO(Derived, Base1, Base2);
/// };
#define DECLARE_TYPEINFO(T, ...)                                                          \
 private:                                                                                 \
    static constexpr P4::RTTI::TypeId static_typeId() { return P4::RTTI::InvalidTypeId; } \
    DECLARE_TYPEINFO_COMMON(T, ##__VA_ARGS__)
// static_typeId is private above in order to hide explicit typeId from the base
// class, as in e.g.
// class A {
//  DECLARE_TYPEINFO_WITH_TYPEID(A, 42);
// };
// class B : public A {
//   DECLARE_TYPEINFO(A, B);
// };
// We need to ensure static_typeId is not visible inside `B`, otherwise
// TypeIdResolver will be confused.

/// Same as DECLARE_TYPEINFO but allows for explicit typeid to be specified.
/// No further checks for typeid collision, etc. are performed. Example:
/// struct Base {
///   ...
///   DECLARE_TYPEINFO_WITH_TYPEID(Base, 42);
/// };
#define DECLARE_TYPEINFO_WITH_TYPEID(T, Id, ...)                                       \
 public:                                                                               \
    static constexpr P4::RTTI::TypeId static_typeId() { return P4::RTTI::TypeId(Id); } \
    DECLARE_TYPEINFO_COMMON(T, ##__VA_ARGS__)

/// Declare typinfo for a given class `T` combining discriminator value and
/// typeid of `InnerT`.
///
/// Use with caution, no checks for typeid clashes, etc. are
/// performed. Discriminator from inner typeid is stripped off.  Examples of
/// usage include `Vector<T>` and `IndexedVector<T>`.
#define DECLARE_TYPEINFO_WITH_DISCRIMINATOR(T, Discriminator, InnerT, ...)                 \
 public:                                                                                   \
    static constexpr P4::RTTI::TypeId static_typeId() {                                    \
        return P4::RTTI::combineTypeIdWithDiscriminator(P4::RTTI::TypeId(Discriminator),   \
                                                        P4::RTTI::TypeInfo<InnerT>::id()); \
    };                                                                                     \
    DECLARE_TYPEINFO_COMMON(T, ##__VA_ARGS__)

// Common typeinfo boilerplate methods. Note that they are marked pure / const, so consecutive
// calls to is<> / to<> on the same pointer could be eliminated by a compiler.
//   - typeId() / isA are const as they do not access any global state, they
//     always return the same value (for same input arguments)
//   - toImpl() is pure. It only access global state via its arguments (this pointer). So for
//     the same `this` the result is also the same.
#define DECLARE_TYPEINFO_COMMON(T, ...)                                                            \
 public:                                                                                           \
    static constexpr T *rttiEnabledMarker(T *);                                                    \
    using TypeInfo = P4::RTTI::TypeInfo<T, ##__VA_ARGS__>;                                         \
    [[nodiscard, gnu::const]] P4::RTTI::TypeId typeId() const noexcept override {                  \
        return TypeInfo::id();                                                                     \
    }                                                                                              \
    [[nodiscard, gnu::const]] bool isA(P4::RTTI::TypeId typeId) const noexcept override {          \
        return TypeInfo::isA(typeId);                                                              \
    }                                                                                              \
                                                                                                   \
 protected:                                                                                        \
    [[nodiscard, gnu::pure]] const void *toImpl(P4::RTTI::TypeId typeId) const noexcept override { \
        return TypeInfo::isA(typeId) ? TypeInfo::dyn_cast(typeId, this) : nullptr;                 \
    }

#endif /* LIB_RTTI_H_ */
